# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import os

from absl import app, flags

import pandas
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

FLAGS = flags.FLAGS


def plot_basis(df, csv_root):
  plt.figure(figsize=(20, 20))
  df['ask_basis_10'].plot(color='b', legend=True)
  df['ask_basis_15'].plot(color='g', legend=True)
  df['ask_basis_30'].plot(color='r', legend=True)
  df['ask_basis_60'].plot(color='y', legend=True)
  df['ask_basis_90'].plot(color='k', legend=True)
  plt.savefig(os.path.join(csv_root, 'ask_basis_plot.png'))

  plt.figure(figsize=(20, 20))
  df['bid_basis_10'].plot(color='b', legend=True)
  df['bid_basis_15'].plot(color='g', legend=True)
  df['bid_basis_30'].plot(color='r', legend=True)
  df['bid_basis_60'].plot(color='y', legend=True)
  df['bid_basis_90'].plot(color='k', legend=True)
  plt.savefig(os.path.join(csv_root, 'bid_basis_plot.png'))

  plt.figure(figsize=(20, 20))
  df['mid_basis_10'].plot(color='b', legend=True)
  df['mid_basis_15'].plot(color='g', legend=True)
  df['mid_basis_30'].plot(color='r', legend=True)
  df['mid_basis_60'].plot(color='y', legend=True)
  df['mid_basis_90'].plot(color='k', legend=True)
  plt.savefig(os.path.join(csv_root, 'mid_basis_plot.png'))


def calculate_basis(df):
  df['ask_diff'] = df['true_ask_x'] - df['true_ask_y']
  df['bid_diff'] = df['true_bid_x'] - df['true_bid_y']
  df['mid_diff'] = (df['true_ask_x'] + df['true_bid_x'] - df['true_ask_y'] - df['true_bid_y']) / 2
  for win_size in [10, 15, 30, 60, 90]:
    df['ask_basis_%d' % win_size] = \
        df['ask_diff'].rolling(window=win_size * 60, min_periods=1).mean()
    df['bid_basis_%d' % win_size] = \
        df['bid_diff'].rolling(window=win_size * 60, min_periods=1).mean()
    df['mid_basis_%d' % win_size] = \
        df['mid_diff'].rolling(window=win_size * 60, min_periods=1).mean()
  return df


def merge_dfs(df1, df2):
  return pandas.merge(df1, df2, how='inner', left_index=True, right_index=True)


def sample_df_per_second(df):
  return df.resample('S', closed='left', label='left').last()


def read_csv_into_df(csv_root, csv):
  csv_dir = os.path.join(csv_root, csv)
  df = pandas.read_csv(csv_dir, sep=',', header=0)
  df['datetime'] = pandas.to_datetime(df['timestamp'], unit='ns')
  df = df[['datetime', 'true_ask', 'true_bid']]
  df = df.set_index('datetime')
  return df


def main(argv):
  csv_root = FLAGS.csv_root
  csv1 = FLAGS.csv1
  csv2 = FLAGS.csv2
  assert csv_root, '--csv_root must be specified.'
  assert csv1, '--csv1 must be specified.'
  assert csv2, '--csv2 must be specified.'

  df1 = read_csv_into_df(csv_root, csv1)
  df2 = read_csv_into_df(csv_root, csv2)
  sampled_df1 = sample_df_per_second(df1)
  sampled_df2 = sample_df_per_second(df2)
  merged_df = merge_dfs(sampled_df1, sampled_df2)
  calculated_df = calculate_basis(merged_df)

  plot_basis(calculated_df, csv_root)

  return 0


if __name__ == '__main__':
  flags.DEFINE_string('csv_root', None, 'Output csv files root directory.')

  flags.DEFINE_string('csv1', None, 'First csv file name.')

  flags.DEFINE_string('csv2', None, 'Second csv file name.')

  app.run(main)
